# coding: UTF-8
import argparse
import os
import time
from importlib import import_module

import numpy as np
import torch
from torch.utils.data import DataLoader

from dataset.CMN_dataset import CMNDataset
from loss.FocalLoss import FocalLoss
from train_eval.train_eval import train, test
from utils.utils import get_time_dif

parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--dataset', type=str, required=True, help='choose a dataset')
parser.add_argument('--model', type=str, required=True, help='choose a model: bert, ERNIE')
parser.add_argument('--mode', type=str, required=True, help='choose a mode: train, test')
args = parser.parse_args()

if __name__ == '__main__':
    dataset = 'data/' + args.dataset
    print('Training on dataset:', dataset)
    model_name = args.model
    x = import_module('models.' + model_name)
    config = (x.Config(dataset))
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True
    
    start_time = time.time()
    print("Loading data...")
    train_dataset = CMNDataset(config, 'train')
    val_dataset = CMNDataset(config, 'val')
    test_dataset = CMNDataset(config, 'test')

    train_dataloader = DataLoader(
        train_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    val_dataloader = DataLoader(
        val_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    test_dataloader = DataLoader(
        test_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)
    model = x.Model(config).to(config.device)
    if args.mode == 'train':
        train(
            config, 
            model, 
            train_dataloader, 
            val_dataloader, 
            test_dataloader, 
            loss_function=FocalLoss(class_num=config.num_classes), 
        )
    if args.mode == 'test':
        test(
            config, 
            model, 
            test_dataloader, 
            loss_function=FocalLoss(class_num=config.num_classes), 
        )
